---
title: "R Notebook"
output: html_notebook
---

# Simulate asymptotic Thompson sampling
```{r}
J = 10 #number of batches

sigma = 1
size = 0.05

batched_TS_asymptotic = function(mu1, mu0){
  
  pi = 0.5 #initialization for first batch
  x1 = 0
  x0 = 0
  q1 = 0
  
  for (j in c(1:J)){
  
    x1 = x1 + rnorm(1, mu1*pi, sqrt(pi)*sigma)
    x0 = x0 + rnorm(1, mu0*(1-pi), sqrt(1-pi)*sigma)
    q1 = q1 + pi
    q0 = j - q1
  
    Delta = (x1/q1) - (x0/q0)
    alpha = sqrt((sigma^2)*j/(q1*(j-q1)))
  
    pi = pnorm(Delta/alpha)
    j = j + 1
  }
  return(c(x1, x0, q1, q0))
} 

#sanity check
bandit_replications = replicate(n = 50000, batched_TS_asymptotic(0,0))
```


# Computing critical values
```{r}
critical_value = function(mu1_h, mu0_h){
  x1 = bandit_replications[1,]
  x0 = bandit_replications[2,]
  q1 = bandit_replications[3,]
  q0 = bandit_replications[4,]
  
  LR_ratio_replications = (mu1_h*x1) + (mu0_h*x0) - 
    (q1*mu1_h^2/2) - (q0*mu0_h^2/2)
  critical_value = quantile(LR_ratio_replications, 1 - size)
  
  return(critical_value)
}
critical_value = Vectorize(critical_value)
```

# Asymptotic power computation using the Girsanov theorem
```{r}
asymptotic_power2 = function(mu1_h, mu0_h){
  c_value = critical_value(mu1_h, mu0_h)
  
  x1 = bandit_replications[1,]
  x0 = bandit_replications[2,]
  q1 = bandit_replications[3,]
  q0 = bandit_replications[4,]
  
  LR_ratio_replications = (mu1_h*x1) + (mu0_h*x0) - (q1*mu1_h^2/2) - (q0*mu0_h^2/2)
  LR_ratio_replications = log(exp(LR_ratio_replications)/mean(exp(LR_ratio_replications)))
  
  power = mean((exp(LR_ratio_replications))*(LR_ratio_replications >= c_value))
  return(power)
}
asymptotic_power2 = Vectorize(asymptotic_power2)
```


# Plotting power envelope
```{r}
library(tidyverse)

c1 = seq(-1.1, 1.1, 0.025)
c2 = seq(-1.1, 1.1, 0.025)
data = expand.grid(C1 = c1, C2 = c2) %>%
        mutate(Z = asymptotic_power2(C1, C2))
```

```{r}
data = data %>% 
        mutate(Z = if_else(C1 == 0 & C2 == 0, 0, Z))

ggplot(data, aes(C1, C2, z = Z)) +
  geom_contour() +
  geom_contour_filled() +
  xlab((expression(paste(eta[1])))) +
  ylab((expression(paste(eta[0])))) +
  theme(#aspect.ratio = 1,
        text = element_text(size=13),
        axis.text = element_text(size=12),
        axis.title = element_text(size=13),
        #legend.position = c(0.85, 0.3)
        )

#setwd("/Users/akarun/Library/CloudStorage/Dropbox/Optimal sequential tests/Figures")
ggsave("Power_envelope_TS.png")
```

# Monte-Carlo simulation - 1
```{r}
TS_power_simulation = function(mu, n){
  pi = 0.5 #initialization for first batch
  x1 = 0
  x0 = 0
  q1 = 0
  
  for (j in c(1:J)){
    n1 = floor(n*pi)
  
    x1 = x1 + n1*(mu/n) + (1/sqrt(n))*sum(sqrt(3)*runif(n1, -1, 1))
    x0 = x0 + (n-n1)*(mu/n) + (1/sqrt(n))*sum(sqrt(3)*runif(n - n1, -1, 1))
    q1 = q1 + (n1/n)
    q0 = j - q1
  
    Delta = (x1/q1) - (x0/q0)
    alpha = sqrt((sigma^2)*j/(q1*(j-q1)))
  
    pi = pnorm(Delta/alpha)
    j = j + 1
  }
  
  test_fn = mu*(x1) + mu*(x0) - (q1*mu^2/2) - (q0*mu^2/2)
  return((test_fn >= critical_value(mu, mu)))
} 

finite_sample_power = function(mu, n){
  replications = replicate(5000, expr = TS_power_simulation(mu, n))
  return(mean(replications))
}

TS_power_simulation = Vectorize(TS_power_simulation)
finite_sample_power = Vectorize(finite_sample_power)
```

```{r}
library(tidyverse)
values = seq(0.1, 1, 0.1)
df1 = data.frame(seq = values)

power_values = df1 %>% 
  mutate(power_n20 = finite_sample_power(seq, 20),
         power_n50 = finite_sample_power(seq, 50),
         power_n100 = finite_sample_power(seq, 100))
```

```{r}
power.df1 = power_values %>%
  mutate(power_infty = asymptotic_power2(seq, seq)) 

names(power.df1) = c("delta_seq", "20", "50", "100", "Asymptotic power")

power.df2 = power.df1 %>% 
  pivot_longer(-delta_seq,
               names_to = "n",
               values_to = "Power") %>%
  mutate(n = factor(n, levels = c("20", "50", "100", "Asymptotic power")))

ggplot(power.df2, aes(x = delta_seq,
                      y = Power,
                      color = n)) +
  geom_point() +
  geom_line() +
  xlab((expression(paste("Scaled outcome means (", eta, ",", eta, ")")))) +
  ylab("Power of Neyman-Pearson tests") +
  theme(#aspect.ratio = 4/3,
        text = element_text(size=13),
        axis.text = element_text(size=12),
        axis.title = element_text(size=13),
        legend.position = c(0.85, 0.22)
        ) +
  scale_color_brewer(palette = "RdGy")

#setwd("/Users/akarun/Library/CloudStorage/Dropbox/Optimal sequential tests/Figures")
ggsave("Power_comparison_TS.png")
```


# Monte-Carlo simulation - 2
```{r}
TS_power_simulation2 = function(mu, n){
  pi = 0.5 #initialization for first batch
  x1 = 0
  x0 = 0
  q1 = 0
  
  for (j in c(1:J)){
    n1 = floor(n*pi)
  
    x1 = x1 + n1*(mu/n) + (1/sqrt(n))*sum(sqrt(3)*runif(n1, -1, 1))
    x0 = x0  + (1/sqrt(n))*sum(sqrt(3)*runif(n - n1, -1, 1))
    q1 = q1 + (n1/n)
    q0 = j - q1
  
    Delta = (x1/q1) - (x0/q0)
    alpha = sqrt((sigma^2)*j/(q1*(j-q1)))
  
    pi = pnorm(Delta/alpha)
    j = j + 1
  }
  
  test_fn = mu*(x1) - (q1*mu^2/2)
  return((test_fn >= critical_value(mu, 0)))
} 

finite_sample_power2 = function(mu, n){
  replications = replicate(5000, expr = TS_power_simulation2(mu, n))
  return(mean(replications))
}

TS_power_simulation2 = Vectorize(TS_power_simulation2)
finite_sample_power2 = Vectorize(finite_sample_power2)
```

```{r}
library(tidyverse)
values = seq(-2, 1.4, 0.1)

power_values = data.frame(seq = values)
power_values = power_values %>% 
  mutate(power_n20 = finite_sample_power2(seq, 20),
         power_n50 = finite_sample_power2(seq, 50),
         power_n100 = finite_sample_power2(seq, 100))

```

```{r}
power.df1 = power_values %>%
  mutate(power_infty = asymptotic_power2(seq, 0))

k = which(power_values$seq == 0)
power.df1$power_n20[k] = NA
power.df1$power_n50[k] = NA
power.df1$power_n100[k] = NA
power.df1$power_infty[k] = NA

names(power.df1) = c("delta_seq", "20", "50", "100", "Asymptotic power")

power.df2 = power.df1 %>% 
  pivot_longer(-delta_seq,
               names_to = "n",
               values_to = "Power") %>%
  mutate(n = factor(n, levels = c("20", "50", "100", "Asymptotic power")))

ggplot(power.df2, aes(x = delta_seq,
                      y = Power,
                      color = n)) +
  geom_point() +
  geom_line() +
  xlab((expression(paste("Scaled outcome means (", eta, ", 0)")))) +
  ylab("Power of Neyman-Pearson tests") +
  theme(#aspect.ratio = 4/3,
        text = element_text(size=13),
        axis.text = element_text(size=12),
        axis.title = element_text(size=13),
        legend.position = c(0.4, 0.78)
        ) +
  scale_color_brewer(palette = "RdGy")

#setwd("/Users/akarun/Library/CloudStorage/Dropbox/Optimal sequential tests/Figures")
ggsave("Power_comparison_TS2.png")
```